# -*- coding: utf-8 -*-

## @package dmiParser
# @author: Ludovic Barriere
# @date: 07/02/2013
# @summary: This module supplies a simple access to Schur complement and assembly
# operators of domains of a Patran/Nastran beam model (../test/beam.dat, ../test/beam.pch).


import scipy.sparse as sp
import scipy.sparse.linalg as spla
from numpy import dot as npdot
from numpy import identity as npid
from numpy import copy as npcopy
from numpy import array as nparray
from numpy import argsort as npargsort
from numpy import size as npsize
import numpy as np
from numpy.linalg import solve as npsolve
from dmiParser import readPunch

NDOF_PER_NODE = 6
DEBUG = False
## class domain
# define a domain and stores its data
class domain():
    
    def __init__(self,Id,KOO=None,KOA=None,KAO=None,KAA=None,PO=None,PA=None):
        
        self.Id = Id
        self.koo = KOO
        self.kaa = KAA
        self.kao = KAO
        self.koa = KOA
        self.po = PO
        self.pa = PA
        self.Saa = None
        self.Baa = None
    
    def getDimension(self):
        
        if self.koo!=None and self.kaa!=None:
            self.dim = [self.koo.get_shape()[0],self.kaa.get_shape()[0]]
            print 'dimensions of domain ',self.Id,' :', self.dim[0],' interior d.o.f. and ', self.dim[1], ' boundary d.o.f.'
        else:
            print 'undefined operators'
        return self.dim
    
    def setOperators(self,KOO=None,KOA=None,KAO=None,KAA=None,PO=None,PA=None):
        """
           Setting the decomposed operators of the linear system that govern the equilibrium of the domain.
        """
        out = [self.koo,self.koa,self.kao,self.kaa,self.po,self.pa]
        
        for i,op in enumerate([KOO,KOA,KAO,KAA,PO,PA]):
            if op != None:
                 out[i]=op
        
        self.koo,self.koa,self.kao,self.kaa,self.po,self.pa = out
        
        return 0
        
    def calcSchur(self,output_matrix_type=sp.csr_matrix):
        """
            Computing the condensed operators on degrees of freedom of the boundary of the domain.
        """
#        KOOinv = spla.factorized(self.koo)
        KOO,KAA,KAO,KOA = self.koo.todense(),self.kaa.todense(),self.kao.todense(),self.koa.todense()
        Saa = KAA - npdot(KAO,npsolve(KOO,KOA))
        
        PO, PA = self.po.todense(),self.pa.todense()
        Baa = PA - npdot(KAO,npsolve(KOO,PO))
        
        self.Saa = output_matrix_type(Saa)
        self.Baa = output_matrix_type(Baa) 
        
        return self.Saa,self.Baa
            
    def getSchur(self):
        return self.Saa, self.Baa
    
class pairBoundary():
    """
        Boundary bewteen two domains.
    """
    def __init__(self,sd1,sd2,BMAP=None):
        
        self.sd = [sd1,sd2]
        if BMAP!=None:
            self.bmap = BMAP
        else:
            self.buildBoundaryMap()
        self.Aopl = self.buildLocalAssemblyOperators() # local assembly operators
        self.Aopg = []                                 # global assembly operators
    
    def buildBoundaryMap(self):
        pass

    def buildLocalAssemblyOperators(self):
        """
            Build the local assembly operators of the local boundary between two domains.
            Meshes are supposed coincident.
        """
        # number of degrees of freedom on boundary
        self.bnodes = npsize(self.bmap,1)
        self.bdof = self.bnodes*NDOF_PER_NODE
        
        # initialize local assembly operators
        Aopl1 = sp.lil_matrix((self.bdof,self.bdof))
        Aopl2 = sp.lil_matrix((self.bdof,self.bdof))    
        
        # local order of boundary nodes
        BMAP = npcopy(self.bmap)
        BMAP[0,:] = npargsort(BMAP[0,:])
        BMAP[1,:] = npargsort(BMAP[1,:])
        
        # fill assembly operators: local numbering of boundary virtual nodes is 0,1,2,... bnodes 
        # Hence, assembly operator maps local boundary node of domain to local boundary virtual node 
        
        for nd in xrange(self.bnodes):
            for d in xrange(NDOF_PER_NODE):
                Aopl1[ nd*NDOF_PER_NODE + d, BMAP[0,nd]*NDOF_PER_NODE + d] = 1 
                Aopl2[ nd*NDOF_PER_NODE + d, BMAP[1,nd]*NDOF_PER_NODE + d] = 1
        
        if DEBUG: 
            print "Aopl",self.sd[0].Id," = ", Aopl1.todense()       
            print "Aopl",self.sd[1].Id," = ", Aopl2.todense()
        
        return [Aopl1,Aopl2]
        
#    def buildAssemblyOperators(self):
#        
#        
#        return [Aopg1,Aopg2]

class boundary():
    """
        Global boundary: assembly of local boundaries (a.k.a. pair boundaries)
    """
    
    
    def __init__(self,domains,pairBoundaries):
        
        self.pb = pairBoundaries
        self.Aopg = {}          # global assembly operators
        self.domains = domains
    
    def buildAssemblyOperators(self):
        """
            Assemble the local assembly operators of the local boundaries. Resulting in the global assembly operators for each domain.
        """
        # calculate total number of boundary nodes:
        self.bnodes = 0
        for Gij in self.pb:
            self.bnodes += Gij.bnodes
            
        # numbering of global boundary is defined arbitrarily according to ordering of
        # pair boundaries.
        self.bdof = self.bnodes * NDOF_PER_NODE
        
        bnode = 0
        
        # 
        for Gij in self.pb:
            sdi = Gij.sd[0].Id
            sdj = Gij.sd[1].Id
            print "Assembling domains ",sdi," and ", sdj
            ## first domain of the local boundary:
            
            # check existence of global assembly operator of the domain:
            try:
                self.Aopg[sdi]
            except KeyError: # if it doesn't exist, initialize it: 
                sdi_bdof = self.domains[sdi].getDimension()[1]
                self.Aopg[sdi] = sp.lil_matrix((self.bdof,sdi_bdof))
                
            # assemble local to global: ndl -> pair boundary relative index; ndg -> index in the entire boundary of domain.
            # bnode is the index of the last global boundary node created by the last pair boundary. 
            for ndl,ndg in enumerate(Gij.bmap[0]):
                self.Aopg[sdi][bnode*NDOF_PER_NODE:(bnode+Gij.bnodes-1)*NDOF_PER_NODE+6,ndg*NDOF_PER_NODE:ndg*NDOF_PER_NODE+6] = Gij.Aopl[0][:,ndl*NDOF_PER_NODE:ndl*NDOF_PER_NODE+6]
            
            ## idem for the second domain of the local boundary:
            try:
                self.Aopg[sdj]
            except KeyError:
                sdj_bdof = self.domains[sdj].getDimension()[1]
                self.Aopg[sdj] = sp.lil_matrix((self.bdof,sdj_bdof))
                
            for ndl,ndg in enumerate(Gij.bmap[1]):
                self.Aopg[sdj][bnode*NDOF_PER_NODE:(bnode+Gij.bnodes-1)*NDOF_PER_NODE+6,ndg*NDOF_PER_NODE:ndg*NDOF_PER_NODE+6] = Gij.Aopl[1][:,ndl*NDOF_PER_NODE:ndl*NDOF_PER_NODE+6] 
            
            # update the number of boundary nodes:        
            bnode+=Gij.bnodes
        
        # check correct number of boundary nodes:
        if bnode != self.bnodes:
            raise "ERROR: global boundary assembly not valid"
        
        return 0
        
        
            

if __name__ == "__main__":
    
    filename = "../test/beam-ddm.pch"
    ndomains = 4
    
    operators = readPunch(filename,output=sp.csr_matrix)
    
    # creating domains
    domains = {}
    for Id in xrange(ndomains):
        KOO = operators["KOO"+str(Id+1)]
        KAA = operators["KAA"+str(Id+1)]
        KAO = operators["KAO"+str(Id+1)]
        KOA = operators["KOA"+str(Id+1)]
        PO = operators["PO"+str(Id+1)]
        PA = operators["PA"+str(Id+1)]
        domains[Id] = domain(Id,KOO,KOA,KAO,KAA,PO,PA)
        domains[Id].calcSchur()
        
    # creating boundaries of domains by pairs:
    BMAP12 = nparray([[0,1,2],[0,1,2]]) 
    G12 = pairBoundary(domains[0],domains[1],BMAP12) 
    BMAP23 = nparray([[3,4,5],[0,1,2]])
    G23 = pairBoundary(domains[1],domains[2],BMAP23)
    BMAP34 = nparray([[3,4,5],[0,1,2]])
    G34 = pairBoundary(domains[2],domains[3],BMAP34)
    
    # creating global boundary:
    G = boundary(domains,[G12,G23,G34])
    G.buildAssemblyOperators()
    
    # checking assembly operators:
    if DEBUG:
        for Id in xrange(ndomains):
            print "Aopg of domain ", Id
            print G.Aopg[Id]
    
    # solving problem:
    print "\n", "-"*60
    print " A vous de jouer pour la résolution du problème... ;-)"
    print "-"*60
    